package cn.xiaopengstack.config;

import cn.xiaopengstack.mcp.model.McpConfigEntity;
import cn.xiaopengstack.persistent.repository.McpConfigRepository;
import com.alibaba.fastjson.JSONObject;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.spec.ClientMcpTransport;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.annotation.Resource;
import org.apache.commons.lang3.BooleanUtils;
import org.springframework.ai.autoconfigure.mcp.client.NamedClientMcpTransport;
import org.springframework.ai.autoconfigure.mcp.client.configurer.McpSyncClientConfigurer;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import java.time.Duration;
import java.util.*;

/**
 * @author jiangyangang
 */
@Configuration
public class McpToolsConfig {

    @Resource
    private McpConfigRepository mcpConfigRepository;

    @Bean
    @Primary
    public List<McpSyncClient> mergeClients(McpSyncClientConfigurer mcpSyncClientConfigurer, List<McpSyncClient> mcpSyncClients){

        ArrayList<McpSyncClient> mcpSyncClientsMerger = new ArrayList<>(mcpSyncClients);

        McpConfigEntity mcpConfigEntity = mcpConfigRepository.queryMcpConfig("mcp_csdn_weixin");
        JSONObject mcpConfigJson = mcpConfigEntity.getMcpConfigJson();
        JSONObject mcpConfigValue = mcpConfigEntity.getMcpConfigValue();
        JSONObject mcpServerJson = mcpConfigJson.getJSONObject("mcpServers");

        for (String mcpAppKey : mcpServerJson.keySet()) {
            if (mcpSyncClients.stream().map(client -> client.getServerInfo().name()).toList().contains(mcpAppKey)) {
                throw new RuntimeException(String.format("DB的Mcp和本地配置MCP有冲突，当前tools:%s", mcpAppKey));
            }
            JSONObject mcpAppsJson = mcpServerJson.getJSONObject(mcpAppKey);

            ServerParameters serverParameters = ServerParameters.builder(mcpAppsJson.getString("command"))
                    .args(mcpAppsJson.getJSONArray("args").toJavaList(String.class))
                    .env(null)
                    .build();
            ClientMcpTransport clientMcpTransport = new StdioClientTransport(serverParameters);

            NamedClientMcpTransport namedTransport = new NamedClientMcpTransport(mcpAppKey, clientMcpTransport);

            McpSchema.Implementation clientInfo = new McpSchema.Implementation(mcpAppKey, "1.0.0");
            McpClient.SyncSpec syncSpec = McpClient.sync(namedTransport.transport()).clientInfo(clientInfo).requestTimeout(Duration.parse("PT" + mcpConfigValue.getString("requestTimeout").toUpperCase()));
            syncSpec = mcpSyncClientConfigurer.configure(namedTransport.name(), syncSpec);
            McpSyncClient syncClient = syncSpec.build();

            // 默认执行初始化客户端
            if (Objects.isNull(mcpConfigValue.getBoolean("initialize")) || BooleanUtils.isTrue(mcpConfigValue.getBoolean("initialize"))) {
                syncClient.initialize();
            }

            mcpSyncClientsMerger.add(syncClient);
        }

        return mcpSyncClientsMerger;
    }

    public static void main(String[] args) {

        Duration parse = Duration.parse("PT360S");
        System.out.println(parse.getSeconds());
    }

    @Bean("syncMcpToolCallbackProvider")
    public SyncMcpToolCallbackProvider syncMcpToolCallbackProvider(@Qualifier("mergeClients") List<McpSyncClient> mergeClients) {
        // 用于记录 name 和其对应的 index
        Map<String, Integer> nameToIndexMap = new HashMap<>();
        // 用于记录重复的 index
        Set<Integer> duplicateIndices = new HashSet<>();

        // 遍历 mcpClients 列表
        for (int i = 0; i < mergeClients.size(); i++) {
            String name = mergeClients.get(i).getServerInfo().name();
            if (nameToIndexMap.containsKey(name)) {
                // 如果 name 已经存在，记录当前 index 为重复
                duplicateIndices.add(i);
            } else {
                // 否则，记录 name 和 index
                nameToIndexMap.put(name, i);
            }
        }

        // 删除重复的元素，从后往前删除以避免影响索引
        List<Integer> sortedIndices = new ArrayList<>(duplicateIndices);
        sortedIndices.sort(Collections.reverseOrder());
        for (int index : sortedIndices) {
            mergeClients.remove(index);
        }

        return new SyncMcpToolCallbackProvider(mergeClients);
    }
}
